import netron
import torch
import os

def visualize_nn(model, shape, path):
    d = torch.rand(shape)
    o = model(d)
    onnx_path = os.path.join(path, f"{model.__class__.__name__}.onnx")
    torch.onnx.export(model, d, onnx_path)
    netron.start(onnx_path)
    

def visualize_RNWN_AUG():
    from NN.ResNet import ResNet18_WordNet_AUG
    model = ResNet18_WordNet_AUG("visualize_RNWN_AUG", 3, 10, "cpu")
    visualize_nn(model, (1, 3, 32, 32), "output")